###############################################################################
## AUTHOR: ALAN YAN
## DATE LAST UPDATED: 07/20/2020
## PURPOSE: PLOT HET EFFECTS RESULTS
###############################################################################
rm(list = ls())

#### NOTES ####

#(1) create a tidier function for random forest output

#### PACKAGES ####
library(pacman)
p_load(tidyverse,
       estimatr,
       broom,
       haven,
       hrbrthemes,
       cjoint,
       formula.tools, 
       DeclareDesign,
       emmeans,
       lmtest,
       cregg,
       grf,
       Cairo)
source("01-yougov-conjoint/03-src/forest-functions.R")

#### FUNCTIONS ####

# ggplot theme
theme_shom_alt <- function (base_size = 11, waffle = FALSE) 
{
  ret <- theme_minimal(base_size = base_size, base_family = "Roboto Condensed") + 
    theme(plot.background = element_rect(fill = "#f5f5f2", 
                                         color = NA), 
          panel.background = element_rect(fill = "#f5f5f2",color = NA), 
          legend.background = element_rect(fill = "#f5f5f2",color = NA))
  if (waffle) {
    ret + theme(axis.text = element_blank())
  }
  else {
    ret
  }
}

#### READ IN DATA ####

results_work <- read_csv("01-yougov-conjoint/01-data/intermediate/hte-forests-work.csv")
results_power <- read_csv("01-yougov-conjoint/01-data/intermediate/hte-forests-power.csv")
results_resp <- read_csv("01-yougov-conjoint/01-data/intermediate/hte-forests-resp.csv")
results_complaints <- read_csv("01-yougov-conjoint/01-data/intermediate/hte-forests-complaints.csv")

#### BIND DATA ####

hte_all <- bind_rows(results_work,results_power,
                     results_resp,results_complaints) %>%
  mutate(outcome = paste("Outcome:\n",outcome,sep = " "),
         outcome = fct_relevel(outcome,
                               "Outcome:\n Prefer to Work",
                               "Outcome:\n More Power",
                               "Outcome:\n More Responsibilities",
                               "Outcome:\n Better Handle Complaints"),
         treatment_short = case_when(
           treatment == "Workers are Shareholders vs\nPrivate Ownership Shareholders" ~ "Workers are Shareholders",
           treatment == "Workers on Board vs\nPrivate Ownership Shareholders" ~ "Workers on Board",
           treatment == "Elect managers vs\nPrivate Ownership Shareholders" ~ "Workers Elect Managers"
         ),
         treatment_short = paste("Treatment:",treatment_short,sep = ' '),
         treatment_short = fct_relevel(treatment_short,
                                       "Treatment: Workers on Board",
                                       "Treatment: Workers are Shareholders",
                                       "Treatment: Workers Elect Managers"))

ate_by_group <- hte_all %>%
  group_by(outcome,treatment_short) %>%
  summarize(cate = unique(cate),
            cate_se = unique(cate_se))
ate_by_group


#### PLOT BY PARTY ID (FIGURE S15) ####

#plot individual-level predictions
hte_all %>%
  drop_na(party_char) %>%
  ggplot(aes(x = id,y = tau.hat)) +
  facet_grid(outcome ~ treatment_short) + 
  geom_linerange(aes(ymin = lower.90,ymax = upper.90,color = party_char),alpha = 0.2) + #CIs around point estimates
  geom_linerange(aes(ymin = lower.95,ymax = upper.95,color = party_char),alpha = 0.01) + 
  geom_hline(data = ate_by_group,aes(yintercept = cate)) + #estimated causal effect full sample
  geom_rect(data = ate_by_group, aes(ymin = cate - 1.96*cate_se,
                  ymax = cate + 1.96*cate_se,
                  xmin = 1, xmax = max(hte_all$id),
                  x = NULL,y = NULL),
            alpha = 0.7) + 
  geom_point(size = 1.5,shape = 21,aes(fill = party_char)) +#estimated unit level causal effect
  scale_fill_manual(values = c("dodgerblue","white","indianred")) + 
  scale_color_manual(values = c("dodgerblue","white","indianred")) + 
  theme_shom_alt(base_size = 16) + 
  theme(axis.text.y = element_blank(),
        plot.caption = element_text(size = 10,hjust = 1),
        plot.caption.position = "plot",
        strip.background = element_rect(fill="grey"),
        strip.text = element_text(color = 'black'),
        legend.position = "bottom",
        strip.text.x = element_text(size = 11),
        strip.text.y = element_text(size = 11)) + 
  labs(x = "",
       y = "Predicted Effect of Treatment vs Private Ownership by Non-Worker Shareholders",
       fill = "Party ID",
       color = "Party ID",
       caption = "Notes: Estimates for the treatment effect for each trial where the comparison is against private ownership by non-worker shareholders are ranked by magnitude and generated via causal forests. 
       The thin vertical line represent the estimated causal effect of each treatment with the vertical shaded regions representing the 95% confidence interval. 
       The horizontal shaded regions represent 95% confidence intervals of individual level treatment effect estimates.") + 
  geom_hline(yintercept = 0,linetype = "dashed",size = 1.5) + 
  coord_flip() + 
  ggsave("01-yougov-conjoint/03-src/output/figures/hte-workplace-dem-pid.pdf",
         dpi = 320,width = 11,height = 14,device = cairo_pdf)

#### PLOT BY PARTY ID (JUST WORK PREFERENCE AND POWER) (FIGURE 5) ####

#plot individual-level predictions
hte_all %>%
  filter(outcome %in% c("Outcome:\n Prefer to Work",
                        "Outcome:\n More Power")) %>%
  drop_na(party_char) %>%
  ggplot(aes(x = id,y = tau.hat)) +
  facet_grid(outcome ~ treatment_short) + 
  geom_linerange(aes(ymin = lower.90,ymax = upper.90,color = party_char),alpha = 0.2) + #CIs around point estimates
  geom_linerange(aes(ymin = lower.95,ymax = upper.95,color = party_char),alpha = 0.01) + 
  geom_hline(data = ate_by_group %>%
               filter(outcome %in% c("Outcome:\n Prefer to Work",
                                     "Outcome:\n More Power")),
             aes(yintercept = cate)) + #estimated causal effect full sample
  geom_rect(data = ate_by_group %>%
              filter(outcome %in% c("Outcome:\n Prefer to Work",
                                    "Outcome:\n More Power")),
            aes(ymin = cate - 1.96*cate_se,
                ymax = cate + 1.96*cate_se,
                xmin = 1, xmax = max(hte_all$id),
                x = NULL,y = NULL),
            alpha = 0.7) + 
  geom_point(size = 1.5,shape = 21,aes(fill = party_char)) +#estimated unit level causal effect
  scale_fill_manual(values = c("dodgerblue","white","indianred")) + 
  scale_color_manual(values = c("dodgerblue","white","indianred")) + 
  theme_shom_alt(base_size = 16) + 
  theme(axis.text.y = element_blank(),
        plot.caption = element_text(size = 10,hjust = 1),
        plot.caption.position = "plot",
        strip.background = element_rect(fill="grey"),
        strip.text = element_text(color = 'black'),
        legend.position = "bottom",
        strip.text.x = element_text(size = 11),
        strip.text.y = element_text(size = 11)) + 
  labs(x = "",
       y = "Predicted Effect of Treatment vs Private Ownership by Non-Worker Shareholders",
       fill = "Party ID",
       color = "Party ID",
       caption = "Notes: Estimates for the treatment effect for each trial where the comparison is against private ownership by non-worker shareholders are ranked by magnitude and generated via causal forests. 
       The thin vertical line represent the estimated causal effect of each treatment with the vertical shaded regions representing the 95% confidence interval. 
       The horizontal shaded regions represent 95% confidence intervals of individual level treatment effect estimates.") + 
  geom_hline(yintercept = 0,linetype = "dashed",size = 1.5) + 
  coord_flip() + 
  ggsave("01-yougov-conjoint/03-src/output/figures/hte-workplace-dem-pid-just-work-power.pdf",
         dpi = 320,width = 11,height = 8.5,device = cairo_pdf)

#### PLOT BY INDUSTRY (FIGURE S8) ####

#plot individual-level predictions
hte_all %>%
  drop_na(industry_char) %>%
  ggplot(aes(x = id,y = tau.hat)) +
  facet_grid(outcome ~ treatment_short) + 
  geom_linerange(aes(ymin = lower.90,ymax = upper.90,color = industry_char),alpha = 0.2) + #CIs around point estimates
  geom_linerange(aes(ymin = lower.95,ymax = upper.95,color = industry_char),alpha = 0.01) + 
  geom_hline(data = ate_by_group,aes(yintercept = cate)) + #estimated causal effect full sample
  geom_rect(data = ate_by_group, aes(ymin = cate - 1.96*cate_se,
                                     ymax = cate + 1.96*cate_se,
                                     xmin = 1, xmax = max(hte_all$id),
                                     x = NULL,y = NULL),
            alpha = 0.7) + 
  geom_point(size = 1.5,shape = 21,aes(fill = industry_char)) +#estimated unit level causal effect
  scale_fill_manual(values = c("dodgerblue","white")) + 
  scale_color_manual(values = c("dodgerblue","white")) + 
  theme_shom_alt(base_size = 16) + 
  theme(axis.text.y = element_blank(),
        plot.caption = element_text(size = 10,hjust = 1),
        plot.caption.position = "plot",
        strip.background = element_rect(fill="grey"),
        strip.text = element_text(color = 'black'),
        legend.position = "bottom",
        strip.text.x = element_text(size = 11),
        strip.text.y = element_text(size = 11)) + 
  labs(x = "",
       y = "Predicted Effect of Treatment vs Private Ownership by Non-Worker Shareholders",
       fill = "Industry",
       color = "Industry",
       caption = "Notes: Estimates for the treatment effect for each trial where the comparison is against private ownership by non-worker shareholders are ranked by magnitude and generated via causal forests. 
       The thin vertical line represent the estimated causal effect of each treatment with the vertical shaded regions representing the 95% confidence interval. 
       The horizontal shaded regions represent 95% confidence intervals of individual level treatment effect estimates.") + 
  geom_hline(yintercept = 0,linetype = "dashed",size = 1.5) + 
  coord_flip() + 
  ggsave("01-yougov-conjoint/03-src/output/figures/hte-workplace-dem-industry.pdf",
         dpi = 320,width = 11,height = 14,device = cairo_pdf)


#### PLOT BY EDUCATION (FIGURE S9) ####

#plot individual-level predictions
hte_all %>%
  drop_na(education_char) %>%
  ggplot(aes(x = id,y = tau.hat)) +
  facet_grid(outcome ~ treatment_short) + 
  geom_linerange(aes(ymin = lower.90,ymax = upper.90,color = education_char),alpha = 0.2) + #CIs around point estimates
  geom_linerange(aes(ymin = lower.95,ymax = upper.95,color = education_char),alpha = 0.01) + 
  geom_hline(data = ate_by_group,aes(yintercept = cate)) + #estimated causal effect full sample
  geom_rect(data = ate_by_group, aes(ymin = cate - 1.96*cate_se,
                                     ymax = cate + 1.96*cate_se,
                                     xmin = 1, xmax = max(hte_all$id),
                                     x = NULL,y = NULL),
            alpha = 0.7) + 
  geom_point(size = 1.5,shape = 21,aes(fill = education_char)) +#estimated unit level causal effect
  scale_fill_manual(values = c("dodgerblue","indianred")) + 
  scale_color_manual(values = c("dodgerblue","indianred")) + 
  theme_shom_alt(base_size = 16) + 
  theme(axis.text.y = element_blank(),
        plot.caption = element_text(size = 10,hjust = 1),
        plot.caption.position = "plot",
        strip.background = element_rect(fill="grey"),
        strip.text = element_text(color = 'black'),
        legend.position = "bottom",
        strip.text.x = element_text(size = 11),
        strip.text.y = element_text(size = 11)) + 
  labs(x = "",
       y = "Predicted Effect of Treatment vs Private Ownership by Non-Worker Shareholders",
       fill = "Education",
       color = "Education",
       caption = "Notes: Estimates for the treatment effect for each trial where the comparison is against private ownership by non-worker shareholders are ranked by magnitude and generated via causal forests. 
       The thin vertical line represent the estimated causal effect of each treatment with the vertical shaded regions representing the 95% confidence interval. 
       The horizontal shaded regions represent 95% confidence intervals of individual level treatment effect estimates.") + 
  geom_hline(yintercept = 0,linetype = "dashed",size = 1.5) + 
  coord_flip() + 
  ggsave("01-yougov-conjoint/03-src/output/figures/hte-workplace-dem-college.pdf",
         dpi = 320,width = 11,height = 14,device = cairo_pdf)